from typing import Union
import numpy as np
import time
from threading import Thread
from math import radians, degrees
from ezgl import Matrix4x4
from ezgl.functions import dispatchmethod
from ezrobot.utils.rpc import RpcClient
# try:
#     from ezfranka.motion import Kinematics, Affine, JointMotion, JointPathMotion, CartPathMotion, CartImpedanceMotion, CartOnlineMotion
# except:
#     pass


class RobotClient:

    def __init__(self):
        self._proxy = RpcClient("http://127.0.0.1:24123/")
        self.velocity_rel = 1
        self.acceleration_rel = 1
        self.move_thread = None
        self.set_ee(Matrix4x4())

    def set_dynamic_rel(self, vel_rel: float, acc_rel: float=None):
        self.velocity_rel = vel_rel
        if acc_rel is not None:
            self.acceleration_rel = acc_rel
        else:
            self.acceleration_rel = vel_rel

    def get_joint(self) -> list:  # rad
        """
        获得 7 自由度关节角度

        Returns:
        - list, 7 自由度关节角度 (rad)
        """
        return self._proxy.get_joint()

    def set_joint_deg(self, joint: Union[np.ndarray, list]):
        joint[:7] = [radians(j) for j in joint[:7]]
        self.set_joint(joint)

    def set_joint(self, joint: Union[np.ndarray, list]):
        self._proxy.set_joint(joint)

    def set_cartesian(self, pose: Matrix4x4):
        """
        设置笛卡尔空间位姿

        Parameters:
        - pose : Matrix4x4, 目标位姿 base_t_ee
        """
        self._proxy.set_cartesian(pose)
        
    def get_tf(self, link: str=None) -> Matrix4x4:
        if link is None:
            link = "panda_hand"
        return self._proxy.get_tf(link)

    def get_pose(self) -> Matrix4x4:
        """
        获取末端执行器位姿
        """
        return self._proxy.get_pose()

    def set_ee(self, hand_t_ee: Matrix4x4):
        """
        设置末端执行器相对于手腕的位姿

        Parameters:
        - hand_t_ee : Matrix4x4
        """
        self._proxy.set_ee(hand_t_ee)

    def get_ee(self) -> Matrix4x4:
        """
        获取末端执行器相对于手腕的位姿 hand_t_ee
        """
        return self._proxy.get_ee()

    # @dispatchmethod
    # def move(self, cart_path_motion: CartPathMotion):
    #     q_init = Affine(self._proxy.get_pose())
    #     ee_t_hand = Affine(self.get_ee().inverse())
    #     cart_path_motion.init(q_init, self.velocity_rel, self.acceleration_rel)

    #     t = 0
    #     duration = cart_path_motion.get_duration()
    #     while t <= duration + 0.01:
    #         target = cart_path_motion.get_q(t)
    #         target = target * ee_t_hand
    #         joint = self.get_joint()
    #         try:
    #             target_joint = Kinematics.analytic_ik(target.matrix(), joint, joint[6])
    #         except Exception as e:
    #             target_joint = Kinematics.inverse(target.vector_rad(), joint, None)

    #         self.set_joint(target_joint)
    #         time.sleep(0.01)
    #         t += 0.01

    # @move.register(JointMotion)
    # @move.register(JointPathMotion)
    # def _(self, joint_motion: Union[JointMotion, JointPathMotion]):
    #     q_init = self.get_joint()
    #     joint_motion.init(q_init, self.velocity_rel, self.acceleration_rel)

    #     t = 0
    #     duration = joint_motion.get_duration()
    #     while t <= duration:
    #         q = joint_motion.get_q(t)
    #         self.set_joint(q)
    #         time.sleep(0.01)
    #         t += 0.01

    # @move.register(CartOnlineMotion)
    # def _(self, cart_online_motion: CartOnlineMotion):
    #     q_init = Affine(self._proxy.get_pose())
    #     ee_t_hand = Affine(self.get_ee().inverse())
    #     cart_online_motion.init(q_init)
    #     cart_online_motion.set_velocity(self.velocity_rel, self.acceleration_rel)

    #     t = 0
    #     while cart_online_motion.is_moving():
    #         target = cart_online_motion.get_pose(t)
    #         target = target * ee_t_hand
    #         joint = self.get_joint()
    #         try:
    #             target_joint = Kinematics.analytic_ik(target.matrix(), joint, joint[6])
    #         except Exception as e:
    #             target_joint = Kinematics.inverse(target.vector_rad(), joint, None)

    #         self.set_joint(target_joint)
    #         time.sleep(0.02)
    #         t += 0.1

    # @move.register(CartImpedanceMotion)
    # def _(self, cart_impedance_motion: CartImpedanceMotion):
    #     q_current = Affine(self._proxy.get_pose())
    #     ee_t_hand = Affine(self.get_ee().inverse())

    #     cart_impedance_motion.init(q_current, self.velocity_rel)

    #     while cart_impedance_motion.is_moving():
    #         q_current = Affine(self._proxy.get_pose())
    #         target = cart_impedance_motion.get_q(q_current) * ee_t_hand
    #         joint = self.get_joint()
    #         try:
    #             target_joint = Kinematics.inverse(target.vector_rad(), joint, None)
    #         except Exception as e:
    #             target_joint = Kinematics.analytic_ik(target.matrix(), joint, joint[6])

    #         self.set_joint(target_joint)
    #         time.sleep(0.025)

    # def move_async(self, motion) -> Thread:
    #     self.move_thread = Thread(target=self.move, args=(motion,), daemon=True)
    #     self.move_thread.start()
    #     return self.move_thread

    # def move_joint(self, joint: Union[np.ndarray, list]):
    #     joint_motion = JointMotion(joint)
    #     self.move(joint_motion)

    # def move_cart(self, tf: Matrix4x4, relative: bool=False):
    #     """
    #     笛卡尔空间运动

    #     Parameters:
    #     - tf : Matrix4x4, 目标位姿 O_T_EE
    #     - relative : bool, optional, default: False, 是否相对于当前位姿运动
    #     """
    #     if relative:
    #         tf_base = self.get_pose()
    #         tf = tf * tf_base

    #     self.move(CartPathMotion(Affine(tf), 0))

    def join(self):
        if self.move_thread is not None:
            self.move_thread.join()

    def toggle_traj(self, visible: bool):
        self._proxy.toggle_traj(visible)

    def add_pose_axis_at_ee(self):
        self._proxy.add_pose_axis_at_ee()

    def add_pose_axis(self, pose: Matrix4x4):
        self._proxy.add_pose_axis(pose)

    def set_pose_axis(self, pose: Matrix4x4):
        self._proxy.set_pose_axis(pose)

    def add_arrow(self, st_pos: np.ndarray, ed_pos: np.ndarray, color: np.ndarray=(1,0,0,1)):
        self._proxy.add_arrow(st_pos, ed_pos, color)

    def set_arrow(self, st_pos: np.ndarray, ed_pos: np.ndarray, color: np.ndarray=(1,0,0,1)):
        self._proxy.set_arrow(st_pos, ed_pos, color)

    def clear_axes(self):
        self._proxy.clear_axes()

    def clear_items(self):
        """
        清除路径和坐标轴
        """
        self.toggle_traj(False)
        self.clear_axes()

    def set_named_axis(self, index, pose: Matrix4x4):
        """
        清除路径和坐标轴
        """
        self._proxy.set_named_axis(index, pose)

    def gripper_move(self, width: float):
        joints = self.get_joint()
        joints.extend([width/2, width/2])
        self.set_joint(joints)
